﻿// Copyright (c) 2014 Morten Bakkedal
// This code is published under the MIT License.

using System;
using System.Collections.Generic;
using System.Diagnostics;

using FuncLib.DualFunctions.Compilation;
using FuncLib.DualFunctions.Expressions;
using FuncLib.Functions;

namespace FuncLib.DualFunctions
{
	public partial class DualFunction
	{
		[DebuggerStepThrough]
		public static DualFunction Exp(DualFunction f)
		{
			return new ExpFunction(f);
		}

		[DebuggerStepThrough]
		public static DualFunction Log(DualFunction f)
		{
			return new LogFunction(f);
		}

		[DebuggerStepThrough]
		public static DualFunction Sqr(DualFunction f)
		{
			return new SqrFunction(f);
		}

		[DebuggerStepThrough]
		public static DualFunction Sqrt(DualFunction f)
		{
			return new SqrtFunction(f);
		}

		[DebuggerStepThrough]
		public static DualFunction Pow(DualFunction f1, DualFunction f2)
		{
			return new PowFunction(f1, f2);
		}

		[DebuggerStepThrough]
		public static DualFunction Pow(DualFunction f, double a)
		{
			// Allow special handling of a few special cases (while not guaranteed to be numerically identical).

			if (a == 1.0)
			{
				return f;
			}

			if (a == 0.5)
			{
				return Sqrt(f);
			}

			if (a == 2.0)
			{
				return Sqr(f);
			}

			return new PowConstantFunction(f, a);
		}

		[DebuggerStepThrough]
		public static DualFunction Pow(double a, DualFunction f)
		{
			throw new NotImplementedException();
		}

		[DebuggerStepThrough]
		public static DualFunction Cos(DualFunction f)
		{
			return new CosFunction(f);
		}

		[DebuggerStepThrough]
		public static DualFunction Sin(DualFunction f)
		{
			return new SinFunction(f);
		}

		[Serializable]
		private class ExpFunction : DualFunction
		{
			private DualFunction f;

			public ExpFunction(DualFunction f)
			{
				this.f = f;

				TestFunction = Function.Exp(f.TestFunction);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc = generator.GetCompiledFunction(f);

				Expression g = Expression.MarkLocal(Expression.Exp(fc.Expression));
				Expression g1 = g;
				Expression g11 = g;

				return generator.AddUnaryFunction(this, fc, g, g1, g11);
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f;
			}
		}

		[Serializable]
		private class LogFunction : DualFunction
		{
			private DualFunction f;

			public LogFunction(DualFunction f)
			{
				this.f = f;

				TestFunction = Function.Log(f.TestFunction);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc = generator.GetCompiledFunction(f);

				Expression fe = Expression.MarkLocal(fc.Expression);
				Expression g = Expression.Log(fe);
				Expression g1 = Expression.MarkLocal(1.0 / fe);
				Expression g11 = Expression.MarkLocal(-g1 / fe);

				return generator.AddUnaryFunction(this, fc, g, g1, g11);
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f;
			}
		}

		[Serializable]
		private class SqrFunction : DualFunction
		{
			private DualFunction f;

			public SqrFunction(DualFunction f)
			{
				this.f = f;

				TestFunction = Function.Sqr(f.TestFunction);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc = generator.GetCompiledFunction(f);

				Expression fe = Expression.MarkLocal(fc.Expression);
				Expression g = fe * fe;
				Expression g1 = 2.0 * fe;
				Expression g11 = 2.0;

				return generator.AddUnaryFunction(this, fc, g, g1, g11);
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f;
			}
		}

		[Serializable]
		private class SqrtFunction : DualFunction
		{
			private DualFunction f;

			public SqrtFunction(DualFunction f)
			{
				this.f = f;

				TestFunction = Function.Sqrt(f.TestFunction);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc = generator.GetCompiledFunction(f);

				Expression fe = Expression.MarkLocal(fc.Expression);
				Expression g = Expression.MarkLocal(Expression.Sqrt(fe));
				Expression g1 = Expression.MarkLocal(0.5 / g);
				Expression g11 = Expression.MarkLocal(-0.5 * g1 / fe);

				return generator.AddUnaryFunction(this, fc, g, g1, g11);
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f;
			}
		}

		[Serializable]
		private class PowFunction : DualFunction
		{
			private DualFunction f1, f2;

			public PowFunction(DualFunction f1, DualFunction f2)
			{
				this.f1 = f1;
				this.f2 = f2;

				TestFunction = Function.Pow(f1.TestFunction, f2.TestFunction);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc1 = generator.GetCompiledFunction(f1);
				GeneratedFunction fc2 = generator.GetCompiledFunction(f2);

				Expression g = Expression.MarkLocal(Expression.Pow(fc1.Expression, fc2.Expression));

				Expression c1 = Expression.MarkLocal(Expression.Pow(fc1.Expression, fc2.Expression - 1.0));
				Expression g1 = fc2.Expression * c1;
				Expression g11 = fc2.Expression * (fc2.Expression - 1.0) * Expression.Pow(fc1.Expression, fc2.Expression - 2.0);

				Expression c2 = Expression.MarkLocal(Expression.Log(fc1.Expression));
				Expression g2 = Expression.MarkLocal(c2 * g);
				Expression g22 = c2 * g2;

				Expression g12 = c1 * (1.0 + c2 * fc2.Expression);

				return generator.AddBinaryFunction(this, fc1, fc2, g, g1, g2, g11, g12, g22);
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f1;
				yield return f2;
			}
		}

		[Serializable]
		private class PowConstantFunction : DualFunction
		{
			private DualFunction f;
			private double a;

			public PowConstantFunction(DualFunction f, double a)
			{
				this.f = f;
				this.a = a;

				TestFunction = Function.Pow(f.TestFunction, a);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc = generator.GetCompiledFunction(f);

				Expression g = Expression.Pow(fc.Expression, a);
				Expression g1 = a * Expression.Pow(fc.Expression, a - 1.0);
				Expression g11 = a * (a - 1.0) * Expression.Pow(fc.Expression, a - 2.0);

				return generator.AddUnaryFunction(this, fc, g, g1, g11);
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f;
			}
		}

		[Serializable]
		private class CosFunction : DualFunction
		{
			private DualFunction f;

			public CosFunction(DualFunction f)
			{
				this.f = f;

				TestFunction = Function.Cos(f.TestFunction);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc = generator.GetCompiledFunction(f);

				Expression fe = Expression.MarkLocal(fc.Expression);
				Expression g = Expression.MarkLocal(Expression.Cos(fe));
				Expression g111 = Expression.MarkLocal(Expression.Sin(fe));
				Expression g1 = -g111;
				Expression g11 = -g;

				return generator.AddUnaryFunction(this, fc, g, g1, g11);
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f;
			}
		}

		[Serializable]
		private class SinFunction : DualFunction
		{
			private DualFunction f;

			public SinFunction(DualFunction f)
			{
				this.f = f;

				TestFunction = Function.Sin(f.TestFunction);
			}

			internal override GeneratedFunction Compile(DualCodeGenerator generator)
			{
				GeneratedFunction fc = generator.GetCompiledFunction(f);

				Expression fe = Expression.MarkLocal(fc.Expression);
				Expression g = Expression.MarkLocal(Expression.Sin(fe));
				Expression g1 = Expression.MarkLocal(Expression.Cos(fe));
				Expression g11 = -g;

				return generator.AddUnaryFunction(this, fc, g, g1, g11);

				// result -> gradient/hessian -> pass that to the next function? yes, build a dependency map (partial order),
				// where functions go out of scope at some point - can also be used to reuse entries in the array
			}

			internal override IEnumerable<DualFunction> GetDependencies()
			{
				yield return f;
			}
		}
	}
}
